/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of LINAL (LINear ALgebra)
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/

#ifndef LINAL_DISPATCHER_HXX
#define LINAL_DISPATCHER_HXX


namespace linal {

  // -------------------------------------------------------------- 
  // ** Dispatcher
  typedef class Device_* Device;
  typedef class Dispatcher_* Dispatcher;

  extern void set_dispatcher( Dispatcher d );
  extern Dispatcher get_dispatcher();
  extern Bool_ comp_dev( Device i, Device j );

  class Dispatcher_ : public Disp_,
                      public Operation_ {
  public:
    typedef class Group_* Group;

    class Group_ : public Disp_ {
    private:
      Int_ _blocksize, _locker, _used;
    public:
      Group_() : _blocksize(0),
                 _locker(0), 
                 _used(0) { }
      Group_(Int_ blocksize, Int_ locker) : _blocksize(blocksize),
                                            _locker(locker),
                                            _used(0) { }
      Group_(const Group_ &b) : _blocksize(b._blocksize),
                                _locker(b._locker),
                                _used(b._used) { }
      virtual~Group_() { }

      Int_ get_blocksize() const { return _blocksize; }
      Int_ get_locker() const { return _locker; }
      Int_ get_used() const { return _used; }

      void acquire_device()       { 
#pragma omp atomic 
        ++_used; 
      }
      void release_device()       { 
#pragma omp atomic
        --_used; 
      }
      Bool_ is_avail() { return _used < _locker; }
      
      virtual Bool_ operator<(const Group_ &b) const  {
        Int_ r_val;

        r_val = this->_blocksize - b._blocksize;
        if (r_val) return (r_val < 0);
        
        r_val = this->_locker - b._locker;
        if (r_val) return (r_val < 0);

        return false;
      }
      virtual Bool_ operator==(const Group_ &b) const {
        return (this->_blocksize == b._blocksize && 
                this->_locker == b._locker);
      }
      virtual Bool_ operator!=(const Group_ &b) const  {
        return !(*this == b);
      }
      virtual std::ostream& content(std::ostream& os) const {
        return os << "(b_mn:: " << _blocksize 
                  << ", locker:: " << _locker 
                  << ", used:: " << _used
                  << ")";
      }
    };
  private:
  protected:

    std::vector< Int_ > _fronts;
    std::vector< Group_ > _groups;
    std::vector< Device > _devices;
    Int_ _locker, _com;

  public:
    Dispatcher_() : _locker(0), _com(0) { }
    virtual ~Dispatcher_() { }

    virtual std::ostream& content(std::ostream& os) const {
      //this->Disp_::content(os);
      os << "- Devices -\n";
      for (Int_ i=0;i<_devices.size();++i)
        os << i << std::endl
           << *(_devices.at(i)) << std::endl;
      
      if (_locker) {
        os << "- Groups -" << std::endl;
        for (Int_ i=0;i<_groups.size();++i)
          os << i 
             << "  " << _groups.at(i)
             << ", (n_dev:: " << (_fronts.at(i+1) - _fronts.at(i)) 
             << ")" << std::endl;
      } else {
        os << "- Groups -" << std::endl;
        for (Int_ i=0;i<_groups.size();++i)
          os << i 
             << "  " << _groups.at(i) 
             << std::endl;
      }
      return os;
    }

    void add_device( Device d ) { _devices.push_back( d ); }
    void add_group( Group_ &group ) { _groups.push_back( group ); }
    void reset() { _fronts.clear(); _groups.clear(); _devices.clear(); _locker = 0; }
    Int_ is_locked() { return _locker; }

    void lock() {
      std::sort( _devices.begin(), _devices.end(), comp_dev );
      std::sort( _groups.begin(), _groups.end() );

      std::unique( _devices.begin(), _devices.end(), comp_dev );
      std::unique( _groups.begin(), _groups.end() );
      
      {
        std::vector< Device >::iterator it = _devices.end();
        --it;
        LINAL_ERROR( ((*it)->get_group() + 1) == _groups.size(), 
                     ">> Mismatch of # of groups to iteration in devices");
      }

      Int_ cur_group = -1;
      for (Int_ i=0;i<_devices.size();++i) {
        Int_ dev_group = _devices.at(i)->get_group();
        if ( cur_group != dev_group ) {
          _fronts.push_back( i );
          cur_group = dev_group;
        }
      }

      _fronts.push_back( _devices.size() );
      _locker = true;
    }
    
    void unlock() { _fronts.clear(); _locker = false; }

    Int_ get_group( FLA_Obj &A ) {
      Int_ 
	r_val = (_groups.size() - 1), 
	blk   = max(A.m, A.n);

      for (Int_ i=0;i<_groups.size();++i) 
	if (_groups[i].get_blocksize() >= blk) {
	  r_val = i;
	  break;
	}
      return r_val;
    }
    Group_* get_group( Int_ group ) { return &_groups[group]; }
    Int_ get_n_device( Int_ group ) { return (_fronts.at(group + 1) - 
                                              _fronts.at(group)); }
    Device get_device( Int_ group ) { return get_device( group, omp_get_thread_num() ); }
    Device get_device( Int_ group, Int_ thread ) {
      Int_ n_dev  = get_n_device(group); 
      return ( n_dev ? 
	       _devices.at( _fronts.at(group) + thread%n_dev ) : NULL );
    }

    void invalidate( Int_ group, Device exe, FLA_Obj &A ) {
      for (Int_ i=0;i<get_n_device(group);++i) {
        Device dev = _devices.at(_fronts.at(group) + i);
        if (dev != exe) {
          dev->invalidate(A);
        }
      }
    }

    void set_computing_model(Int_ com) { _com = com; }
    Int_ get_computing_model() { return _com; }
    
    void stat_begin() {
      for (Int_ i=0;i<_devices.size();++i)
        _devices.at(i)->stat_begin();
    }
    void stat_end() {
      for (Int_ i=0;i<_devices.size();++i)
        _devices.at(i)->stat_end();
    }

    virtual std::ostream& statistics(std::ostream& os) const {
      for (Int_ i=0;i<_devices.size();++i) 
        _devices.at(i)->statistics(os);
      return os;
    }

    virtual Int_ chol_nopiv  ( Int_ uplo, FLA_Obj &A );
    virtual Int_ chol_piv    ( Int_ uplo, FLA_Obj &A, FLA_Obj &p );
    virtual Int_ ldl_nopiv   ( Int_ uplo, FLA_Obj &A );
    virtual Int_ ldl_piv     ( Int_ uplo, FLA_Obj &A, FLA_Obj &p );
    virtual Int_ lu_nopiv    ( FLA_Obj &A );
    virtual Int_ lu_piv      ( FLA_Obj &A, FLA_Obj &p );
    virtual Int_ qr          ( FLA_Obj &A, FLA_Obj &T );


    virtual Int_ gemm        ( Int_ transa, Int_ transb,
                               FLA_Obj &alpha,  FLA_Obj &A,  FLA_Obj &B,
                               FLA_Obj &beta,   FLA_Obj &C );
    
    virtual Int_ herk        ( Int_ uplo, Int_ trans,
                               FLA_Obj &alpha,  FLA_Obj &A,
                               FLA_Obj &beta,   FLA_Obj &C );

    virtual Int_ trmm        ( Int_ side, Int_ uplo, Int_ trans, Int_ diag,
                               FLA_Obj &alpha,  FLA_Obj &A,  FLA_Obj &B );
    
    virtual Int_ trsm        ( Int_ side, Int_ uplo, Int_ trans, Int_ diag,
                               FLA_Obj &alpha,  FLA_Obj &A,  FLA_Obj &B );

    virtual Int_ syrk        ( Int_ uplo, Int_ trans,
                               FLA_Obj &alpha,  FLA_Obj &A,
                               FLA_Obj &beta,   FLA_Obj &C );

  };


}

#endif
